package message

import (
	"time"

	"github.com/Connor1996/badger"
	"github.com/pingcap-incubator/tinykv/proto/pkg/raft_cmdpb"
)

type Callback struct {
	Response *raft_cmdpb.RaftCmdResponse
	Txn      *badger.Txn // used for GetSnap
	done     chan struct{}
}

func (callback *Callback) Done(response *raft_cmdpb.RaftCmdResponse) {
	if callback == nil {
		return
	}
	if response != nil {
		callback.Response = response
	}
	callback.done <- struct{}{}
}

func (callback *Callback) Wait() *raft_cmdpb.RaftCmdResponse {
	select {
	case <-callback.done:
		return callback.Response
	}
}

func (callback *Callback) WaitByTimeout(timeout time.Duration) *raft_cmdpb.RaftCmdResponse {
	select {
	case <-callback.done:
		return callback.Response
	case <-time.After(timeout):
		return callback.Response
	}
}

func NewCallback() *Callback {
	done := make(chan struct{}, 1)
	cb := &Callback{done: done}
	return cb
}
